;;; guile-openai --- An OpenAI API client for Guile
;;; Copyright © 2023 Andrew Whatson <whatson@tailcall.au>
;;;
;;; This file is part of guile-openai.
;;;
;;; guile-openai is free software: you can redistribute it and/or modify
;;; it under the terms of the GNU Affero General Public License as
;;; published by the Free Software Foundation, either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; guile-openai is distributed in the hope that it will be useful, but
;;; WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; Affero General Public License for more details.
;;;
;;; You should have received a copy of the GNU Affero General Public
;;; License along with guile-openai.  If not, see
;;; <https://www.gnu.org/licenses/>.

(define-module (openai completion)
  #:use-module (openai api completion)
  #:use-module (openai client)
  #:use-module (openai utils colorized)
  #:use-module (ice-9 match)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-9 gnu)
  #:use-module (srfi srfi-41)
  #:export (openai-default-completion-model
            openai-default-completion-temperature
            openai-default-completion-top-p

            completion?
            completion-content
            completion-stream

            openai-completion))

(define-once openai-default-completion-model
  (make-parameter 'ada))

(define-once openai-default-completion-temperature
  (make-parameter *unspecified*))

(define-once openai-default-completion-top-p
  (make-parameter *unspecified*))

(define-once openai-default-completion-stream?
  (make-parameter #t))

(define-record-type <Completion>
  (%make-completion content stream)
  completion?
  (content %completion-content)
  (stream completion-stream))

(define (completion-content completion)
  (force (%completion-content completion)))

(define (make-completions result n)
  (define make-one-completion
    (if (stream? result)
        ;; streamed response will be lazily parsed & merged
        (lambda (ix)
          (let* ((strm (responses->content-stream result ix))
                 (cont (delay (string-concatenate (stream->list strm)))))
            (%make-completion cont strm)))
        ;; regular response will be lazily parsed
        (lambda (ix)
          (let* ((cont (response->content result ix))
                 (strm (stream (force cont))))
            (%make-completion cont strm)))))
  (apply values (map make-one-completion (iota n))))

(define (response->content response n)
  (completion-choice-text
   (list-ref (completion-response-choices response) n)))

(define (responses->content-stream responses n)
  (stream-let loop ((responses responses))
    (if (stream-null? responses)
        stream-null
        (let* ((response (stream-car responses))
               (choice   (car (completion-response-choices response)))
               (index    (completion-choice-index choice))
               (content  (completion-choice-text choice)))
          (cond ((not (eqv? index n))        ;; ignore wrong index
                 (loop (stream-cdr responses)))
                ((unspecified? content)      ;; ignore unspecified content
                 (loop (stream-cdr responses)))
                (else
                 (stream-cons content (loop (stream-cdr responses)))))))))

(define (print-completion completion port)
  (newline port)
  (stream-for-each (lambda (content)
                     (display content port))
                   (completion-stream completion)))

(set-record-type-printer! <Completion> print-completion)

(add-color-scheme! `((,completion? completion ,color-stream (GREEN BOLD))))

(define* (openai-completion prompt #:key
                            (model             (openai-default-completion-model))
                            (suffix            *unspecified*)
                            (max-tokens        *unspecified*)
                            (temperature       (openai-default-completion-temperature))
                            (top-p             (openai-default-completion-top-p))
                            (n                 *unspecified*)
                            (stream?           (openai-default-completion-stream?))
                            (logprobs          *unspecified*)
                            (echo              *unspecified*)
                            (stop              *unspecified*)
                            (presence-penalty  *unspecified*)
                            (frequency-penalty *unspecified*)
                            (best-of           *unspecified*)
                            (logit-bias        *unspecified*)
                            (user              (openai-default-user)))
  (let* ((model (if (symbol? model) (symbol->string model) model))
         (stream? (or stream? *unspecified*))
         (request (make-completion-request model prompt suffix max-tokens
                                           temperature top-p n stream? logprobs echo
                                           stop presence-penalty frequency-penalty
                                           best-of logit-bias user))
         (response (send-completion-request request)))
    (make-completions response (if (unspecified? n) 1 n))))
